package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"

	"github.com/zeromicro/go-zero/core/logx"
)

// AES加密
func AESEncrypt(keyStr, dataStr string) (string, error) {
	// Base64解码密钥
	key, err := base64.StdEncoding.DecodeString(keyStr)
	if err != nil {
		return "", err
	}

	// 创建AES加密块
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	// 填充数据
	data := []byte(dataStr)
	data = pkcs7Padding(data, block.BlockSize())

	// 创建加密器
	ciphertext := make([]byte, len(data))
	blockMode := cipher.NewCBCEncrypter(block, key[:block.BlockSize()])
	blockMode.CryptBlocks(ciphertext, data)

	// Base64编码加密后的数据
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// PKCS7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padtext...)
}

// HMAC-MD5签名
func HMACMD5(keyStr, dataStr string) (string, error) {
	// Base64解码密钥
	key, err := base64.RawURLEncoding.DecodeString(keyStr)
	if err != nil {
		logx.Errorf("HMACMD5,解码密钥错误为：%s", err.Error())
		return "", err
	}

	// 创建HMAC-MD5签名器
	h := hmac.New(md5.New, key)
	h.Write([]byte(dataStr))

	// 计算签名
	signature := h.Sum(nil)

	// Base64编码签名
	return base64.StdEncoding.EncodeToString(signature), nil
}

// AES解密
func AESDecrypt(keyStr, encryptedDataStr string) (string, error) {
	// Base64解码密钥
	key, err := base64.StdEncoding.DecodeString(keyStr)
	if err != nil {
		return "", err
	}

	// Base64解码加密数据
	encryptedData, err := base64.StdEncoding.DecodeString(encryptedDataStr)
	if err != nil {
		return "", err
	}

	// 创建AES解密块
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	// 创建解密器
	blockMode := cipher.NewCBCDecrypter(block, key[:block.BlockSize()])
	decrypted := make([]byte, len(encryptedData))
	blockMode.CryptBlocks(decrypted, encryptedData)

	// 去除填充
	decrypted = pkcs7UnPadding(decrypted)

	return string(decrypted), nil
}

// PKCS7去除填充
func pkcs7UnPadding(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	return data[:(length - unpadding)]
}

func MD5(v string) string {
	d := []byte(v)
	m := md5.New()
	m.Write(d)
	return hex.EncodeToString(m.Sum(nil))
}
